package com.bungarus.group.auth;

import org.springframework.security.authentication.AbstractAuthenticationToken;
import org.springframework.security.core.GrantedAuthority;
import org.springframework.util.StringUtils;

import javax.crypto.Cipher;
import javax.crypto.spec.IvParameterSpec;
import javax.crypto.spec.SecretKeySpec;
import javax.security.auth.Subject;
import java.security.GeneralSecurityException;
import java.security.Key;
import java.security.KeyFactory;
import java.security.PublicKey;
import java.security.spec.X509EncodedKeySpec;
import java.util.*;

/**
 * Created by tang on 2019/7/5.
 */
public class SessionToken extends AbstractAuthenticationToken {
    public static final String USER_ID = "userid";
    public static final String TENANT_ID = "tenantid";
    public static final String TIMESTAMP = "timestamp";
    public static final String LIFETIME = "lifetime";

    private static final String ALGORITHM_NAME = "RSA" ;
    private static final String[] ID_NAMES = {USER_ID, TENANT_ID, TIMESTAMP, LIFETIME};

    private Object principal;
    private Object credential;
    private Long userId;
    private Long tenantId;
    private Long timestamp;
    private Long lifetime;

    public SessionToken(Long userId, Long tenantId, Long timestamp, Long lifetime) {
        super(null);
        this.userId = userId;
        this.tenantId = tenantId;
        this.timestamp = timestamp;
        this.lifetime = lifetime;
        super.setAuthenticated(false);
    }

    public SessionToken(Long userId, Long tenantId, Long timestamp, Long lifetime, Collection<? extends GrantedAuthority> authorities) {
        super(authorities);
        this.userId = userId;
        this.tenantId = tenantId;
        this.timestamp = timestamp;
        this.lifetime = lifetime;
        super.setAuthenticated(true);
    }

    public static Map<String, Long> parseIdentities(String token, String key) {
        if(StringUtils.isEmpty(key)) {
            throw new IllegalArgumentException("key is null.");
        }
        Map<String, Long> res = new HashMap<>();
        try {
            StringTokenizer stringTokenizer = new StringTokenizer(token, ".");
            String encryptedData = stringTokenizer.nextToken();
            String encryptedKeyIV = stringTokenizer.nextToken();

            String plainKeyIV = rsaDecrypt(Base64.getDecoder().decode(encryptedKeyIV), toPublicKey(key));
            stringTokenizer = new StringTokenizer(plainKeyIV, ".");
            String aesKey = stringTokenizer.nextToken();
            String aesIV = stringTokenizer.nextToken();

            String plainToken = aesDecrypt(Base64.getDecoder().decode(encryptedData), aesKey, aesIV);
            //pattern: userid:tenantid:timestamp:lifetime
            stringTokenizer = new StringTokenizer(plainToken, ":");
            int i = 0;
            while (stringTokenizer.hasMoreElements()) {
                Long value = Long.valueOf((String)stringTokenizer.nextElement());
                res.put(ID_NAMES[i], value);
                i++;
            }
        }
        catch (Exception e) {
            e.printStackTrace();
        }
        return res;
    }

    private static String rsaDecrypt(byte[] encryptedMessage, Key publicKey) throws Exception {
        Cipher c = Cipher.getInstance(ALGORITHM_NAME);//Cipher.getInstance(ALGORITHM_NAME + "/" + MODE_OF_OPERATION + "/" + PADDING_SCHEME) ;
        c.init(Cipher.DECRYPT_MODE, publicKey);
        byte[] plainText = c.doFinal(encryptedMessage);

        return new String(plainText, "utf-8") ;
    }

    private static String aesDecrypt(byte[] encryptedMessage, String key, String iv) throws Exception {
        try {
            byte[] raw = key.getBytes("utf-8");
            SecretKeySpec skeySpec = new SecretKeySpec(raw, "AES");
            Cipher cipher = Cipher.getInstance("AES/CBC/PKCS5Padding");
            IvParameterSpec ips = new IvParameterSpec(iv.getBytes("utf-8"));
            cipher.init(Cipher.DECRYPT_MODE, skeySpec, ips);
            try {
                byte[] original = cipher.doFinal(encryptedMessage);
                String originalString = new String(original, "utf-8");
                return originalString;
            } catch (Exception e) {
                System.out.println(e.toString());
                return null;
            }
        } catch (Exception ex) {
            System.out.println(ex.toString());
            return null;
        }
    }

    private static PublicKey toPublicKey(String keyStr) throws GeneralSecurityException {
        byte[] keyData = Base64.getDecoder().decode(keyStr);
        X509EncodedKeySpec spec = new X509EncodedKeySpec(keyData);
        KeyFactory fact = KeyFactory.getInstance(ALGORITHM_NAME);
        return fact.generatePublic(spec);
    }

    @Override
    public Object getCredentials() {
        return this.credential;
    }

    public void setCredentials(Object credential) {
        this.credential = credential;
    }

    @Override
    public Object getPrincipal() {
        return tenantId + ":" + userId;
    }

    public Long getUserId() {
        return userId;
    }

    public Long getTenantId() {
        return tenantId;
    }

    public Long getTimestamp() {
        return timestamp;
    }

    public Long getLifetime() {
        return lifetime;
    }
}
